# -*- coding: utf-8 -*-
"""
Created on Wed Nov 18 09:22:00 2020
策略为三种：
一，直接选保真度最高的
二，如果保真度不比之前的高，就选次优的
三，如果保真度不比之前的高，就选最差的
两个比特各自可以选择不同的动作，但每个比特上，选用的动作都是单方向的。
@author: Waikikilick
"""

import numpy as np
from scipy.linalg import expm
from time import *
import multiprocessing as mp
import copy
np.random.seed(1)
T = 10*np.pi
dt = np.pi/4
step_max = T/dt

action_space = np.array([[1,0,0], #可以选择的动作范围，各列的每项分别代表着 sigma x, y, z 前面的系数。
                       [2,0,0], #每次执行的动作都是单独的绕 x, y, z 轴一定角度的旋转
                       [0,1,0], # x, y 方向的值可以取负，但 z 方向的只能取正值
                       [0,2,0],
                       [0,0,1],
                       [0,0,2],
                       [-1,0,0],
                       [-2,0,0],
                       [0,-1,0],
                       [0,-2,0],
                       [0,0,0]])

action_unit_list = []
for action1 in action_space:
    for action2 in action_space:
        action_unit_list.append([action1,action2])
        

s_x = np.mat([[0,1],[1,0]],dtype=complex)
s_y = np.mat([[0,-1j],[1j,0]],dtype=complex)
s_z = np.mat([[1,0],[0,-1]],dtype=complex)
H_c = np.kron((s_x+1j*s_y)/2, (s_x-1j*s_y)/2) + np.kron((s_x-1j*s_y)/2, (s_x+1j*s_y)/2)

init_set_num = 512
target_set_num = init_set_num

qubit_1 = np.mat([[1,0],[0,0]],dtype=complex)
qubit_2 = np.mat([[0,0],[0,1]],dtype=complex)

noise_normal = np.array([ 1.62434536, -0.61175641, -0.52817175, -1.07296862,  0.86540763,
                         -2.3015387 ,  1.74481176, -0.7612069 ,  0.3190391 , -0.24937038,
                          1.46210794, -2.06014071, -0.3224172 , -0.38405435,  1.13376944,
                         -1.09989127, -0.17242821, -0.87785842,  0.04221375,  0.58281521,
                         -1.10061918,  1.14472371,  0.90159072,  0.50249434,  0.90085595,
                         -0.68372786, -0.12289023, -0.93576943, -0.26788808,  0.53035547,
                         -0.69166075, -0.39675353, -0.6871727 , -0.84520564, -0.67124613,
                         -0.0126646 , -1.11731035,  0.2344157 ,  1.65980218,  0.74204416,
                         -0.19183555,-0.88762896,  -0.74715829,  1.6924546 ,  0.05080775,
                         -0.63699565,  0.19091548,  2.10025514,  0.12015895,  0.61720311,
                         0.30017032, -0.35224985, -1.1425182,  -0.34934272,  -0.20889423,
                         0.58662319,  0.83898341,  0.93110208,  0.28558733,  0.88514116, 
                         -0.75439794,  1.25286816,  0.51292982, -0.29809284,  0.48851815,
                         -0.07557171,  1.13162939,  1.51981682,  2.18557541, -1.39649634,
                         -1.44411381, -0.50446586,  0.16003707,  0.87616892,  0.31563495,
                         -2.02220122, -0.30620401,  0.82797464,  0.23009474,  0.76201118, 
                         -0.22232814, -0.20075807,  0.18656139,  0.41005165,  0.19829972,
                         0.11900865, -0.67066229,  0.37756379,  0.12182127,  1.12948391,
                         1.19891788,  0.18515642, -0.37528495, -0.63873041,  0.42349435,
                         0.07734007, -0.34385368,  0.04359686, -0.62000084,  0.69803203])
#noise_normal 为均值为 0 ，标准差为 1 的正态分布的随机数组成的数组
#该随机数由 np.random.seed(1) 生成: np.random.seed(1) \ noise_uniform = np.random.normal(loc=0.0, scale=1.0, size=41)
noise_drift = np.array([1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.,1.])

alpha_num = 8
theta = [np.pi/8,np.pi/4,3*np.pi/8]
theta_1 = theta
theta_2 = theta
theta_3 = theta

alpha = np.linspace(0,np.pi*2,alpha_num,endpoint=False)
alpha_1 = alpha
alpha_2 = alpha
alpha_3 = alpha
alpha_4 = alpha

a_list_complex = np.matrix([[0,0,0,0]],dtype=complex) #第一行用来占位，否则无法和其他行并在一起，在最后要注意去掉这一行
for ii in range(3): #theta_1
    for jj in range(3): #theta_2
        for kk in range(3): #theta_3
            for mm in range(alpha_num): #alpha_1
                for nn in range(alpha_num): #alpha_2
                    for oo in range(alpha_num): #alpha_3
                        for pp in range(alpha_num): #alpha_4
                            
                            a_1_mo = np.cos(theta_1[ii])
                            a_2_mo = np.sin(theta_1[ii])*np.cos(theta_2[jj])
                            a_3_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.cos(theta_3[kk])
                            a_4_mo = np.sin(theta_1[ii])*np.sin(theta_2[jj])*np.sin(theta_3[kk])
                            
                            a_1_real = a_1_mo*np.cos(alpha_1[mm])
                            a_1_imag = a_1_mo*np.sin(alpha_1[mm])
                            a_2_real = a_2_mo*np.cos(alpha_2[nn])
                            a_2_imag = a_2_mo*np.sin(alpha_2[nn])
                            a_3_real = a_3_mo*np.cos(alpha_3[oo])
                            a_3_imag = a_3_mo*np.sin(alpha_3[oo])
                            a_4_real = a_4_mo*np.cos(alpha_4[pp])
                            a_4_imag = a_4_mo*np.sin(alpha_4[pp])
                            
                            a_1_complex = a_1_real + a_1_imag*1j
                            a_2_complex = a_2_real + a_2_imag*1j
                            a_3_complex = a_3_real + a_3_imag*1j
                            a_4_complex = a_4_real + a_4_imag*1j
                            
                            a_complex = np.matrix([[ a_1_complex, a_2_complex, a_3_complex, a_4_complex]])
                            a_list_complex = np.row_stack((a_list_complex,a_complex))
                            
psi_set = np.array(np.delete(a_list_complex,0,axis=0)) # 删除矩阵的第一行
np.random.shuffle(psi_set) #打乱顺序
init_set = psi_set[:init_set_num]


def target_set():
    target_set = psi_set[init_set_num : init_set_num + target_set_num]
    return target_set

target_set = target_set()

def step(psi,target_psi,action_unit,noise_unit):
    action1, action2 = action_unit[0], action_unit[1]
    H_12 = (np.kron(qubit_1,(1+noise_unit[0])*action1[0]*s_x) + np.kron(qubit_1,(1+noise_unit[0])*action1[1]*s_y) + np.kron(qubit_1,(1+noise_unit[0])*action1[2]*s_z)+ np.kron(qubit_2,(1+noise_unit[1])*action2[0]*s_x) + np.kron(qubit_2,(1+noise_unit[1])*action2[1]*s_y) + np.kron(qubit_2,(1+noise_unit[1])*action2[2]*s_z))/2
    
    H = H_12 + H_c
    U = expm(-1j * H * dt) 
    psi_ = U * psi
    fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
    return fid, psi_

#动作直接选最优的
def step0(psi,target_psi,F):
    fid_list = []
    psi_list = []
    
    for action1 in range(len(action_space)): #action1 是第一个比特上采用的动作
        for action2 in range(len(action_space)): #action2 是第二个比特上采用的动作
        
            H_12 = (np.kron(qubit_1,action_space[action1,0]*s_x) + np.kron(qubit_1,action_space[action1,1]*s_y) + np.kron(qubit_1,action_space[action1,2]*s_z)+ np.kron(qubit_2,action_space[action2,0]*s_x) + np.kron(qubit_2,action_space[action2,1]*s_y) + np.kron(qubit_2,action_space[action2,2]*s_z))/2
    
            H = H_12 + H_c
            U = expm(-1j * H * dt) 
            psi_ = U * psi
            fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
            
            psi_list.append(psi_)
            fid_list.append(fid)
    best_action = fid_list.index(max(fid_list))
    
    best_fid = max(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_


#动作选最优的，或者最差的
def step1(psi,target_psi,F):
    fid_list = []
    psi_list = []
    
    for action1 in range(len(action_space)): #action1 是第一个比特上采用的动作
        for action2 in range(len(action_space)): #action2 是第二个比特上采用的动作
        
            H_12 = (np.kron(qubit_1,action_space[action1,0]*s_x) + np.kron(qubit_1,action_space[action1,1]*s_y) + np.kron(qubit_1,action_space[action1,2]*s_z)+np.kron(qubit_2,action_space[action2,0]*s_x) + np.kron(qubit_2,action_space[action2,1]*s_y) + np.kron(qubit_2,action_space[action2,2]*s_z))/2
    
            H = H_12 + H_c
            U = expm(-1j * H * dt) 
            psi_ = U * psi
            fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
            
            psi_list.append(psi_)
            fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        
        best_action = fid_list.index(min(fid_list))
        best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    return best_action, best_fid, psi_

#动作选最优的，或者次优的
def step2(psi,target_psi,F):
    fid_list = []
    psi_list = []
    for action1 in range(len(action_space)): #action1 是第一个比特上采用的动作
        for action2 in range(len(action_space)): #action2 是第二个比特上采用的动作
        
            H_12 = (np.kron(qubit_1,action_space[action1,0]*s_x) + np.kron(qubit_1,action_space[action1,1]*s_y) + np.kron(qubit_1,action_space[action1,2]*s_z)+np.kron(qubit_2,action_space[action2,0]*s_x) + np.kron(qubit_2,action_space[action2,1]*s_y) + np.kron(qubit_2,action_space[action2,2]*s_z))/2
    
            H = H_12 + H_c
            U = expm(-1j * H * dt) 
            psi_ = U * psi
            fid = (np.abs(psi_.H * target_psi) ** 2).item(0).real 
            
            psi_list.append(psi_)
            fid_list.append(fid)
        
    if F < max(fid_list):
        best_action = fid_list.index(max(fid_list))
        best_fid = max(fid_list)
    else:
        psi_list_ = copy.deepcopy(psi_list)
        fid_list_ = copy.deepcopy(fid_list)
        
        del psi_list_[fid_list_.index(max(fid_list_))]
        del fid_list_[fid_list_.index(max(fid_list_))]
        
        best_action = fid_list.index(max(fid_list_))
        
        best_fid = max(fid_list_)
        
        # best_action = fid_list.index(min(fid_list))
        # best_fid = min(fid_list)
    psi_ = psi_list[best_action]
    # print(best_action)
    return best_action, best_fid, psi_
#---------------------------------------------------------------------------------
#--------------------------------------------------------------------------------
def job(target_psi):
    
    fid_list_noise = [] # 中间元素为列表，每个列表元素为[a,b,c,d,...],其中 a 为无噪声时， 其他为有噪声时
    target_psi = np.mat(target_psi).T
    for k in range(init_set_num):
        psi1 = np.mat(init_set[k]).T
        fid_list = [] # 元素为[a,b,c,d,...],其中 a 为无噪声时， 其他为有噪声时
        
        action_list_1 = []
        action_list_2 = []
        action_list_0 = []
        
        psi = psi1
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        
        fid_max1 = F
        fid_max2 = F
        fid_max0 = F
        
        fid_list_1 = [F]
        fid_list_2 = [F]
        fid_list_0 = [F]
        
        step_n = 0
        while True:
            action, F, psi_ = step1(psi,target_psi,F)
            action_list_1.append(action)
            fid_list_1.append(F)
            fid_max1 = max(F,fid_max1)
            psi = psi_
            step_n += 1
            if fid_max1>0.999 or step_n>step_max:
                break
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step2(psi,target_psi,F)
            action_list_2.append(action)
            fid_list_2.append(F)
            fid_max2 = max(F,fid_max2)
            psi = psi_
            step_n += 1
            if fid_max2>0.999 or step_n>step_max:
                break 
            
        step_n = 0
        F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
        psi = psi1
        while True:
            action, F, psi_ = step0(psi,target_psi,F)
            action_list_0.append(action)
            fid_list_0.append(F)
            fid_max0 = max(F,fid_max0)
            psi = psi_
            step_n += 1
            if fid_max0>0.999 or step_n>step_max:
                break 
            
        fid_max = max(fid_max0,fid_max1,fid_max2) 
        fid_list.append(fid_max)
        fid_max_index = [fid_max1,fid_max2,fid_max0].index(max([fid_max1,fid_max2,fid_max0]))
        
        if fid_max_index == 0: 
            actions_list = action_list_1
            fids_list = fid_list_1
        elif fid_max_index == 1:
            actions_list = action_list_2
            fids_list = fid_list_2
        elif fid_max_index == 2:
            actions_list = action_list_0
            fids_list = fid_list_0
        
        max_index = fids_list.index(fid_max)
        actions_list = actions_list[0:max_index+1]
        
        #在对应于最大保真度的动作策略上施加噪声
         
        for noise_am in noise_am_list:
            
            noise_list = noise_am * noise_drift 
            # noise_list = noise_am * noise_normal
            psi = psi1
            F = (np.abs(psi1.H * target_psi) ** 2).item(0).real 
            fid_noise_list = [F]
            i = 0
            for action in actions_list:
                action_unit = action_unit_list[action]
                noise_unit = [noise_list[i],noise_list[i+1]]
                F, psi_ =  step(psi, target_psi, action_unit, noise_unit)
                i += 2
                fid_noise_list.append(F)
                psi = psi_
            fid_noise = fid_noise_list[max_index]
            fid_list.append(fid_noise)
            
        fid_list_noise.append(fid_list)
        
    fid_list_noise = np.array(fid_list_noise)
            
    return np.mean(fid_list_noise,axis=0)

# noise_am_list = [0.00,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.10]
noise_am_list = [-0.020,-0.018,-0.016,-0.014,-0.012,-0.010,-0.008,-0.006,-0.004,-0.002,0.000,0.002,0.004,0.006,0.008,0.010,0.012,0.014,0.016,0.018,0.020]

def multicore():
    pool = mp.Pool()
    F_list_noise = pool.map(job, target_set)
    return np.mean(np.array(F_list_noise),axis=0)

    
if __name__ == '__main__':

    
    time1 = time()
    
    F_list_noise =  multicore()
    
    print('\n F_list_noise = \n',F_list_noise) 
    
    time2 = time()
    
    print('\n time_cost is: ',time2-time1)
    
    
# normal
# noise_am_list = [0.00,0.01,0.02,0.03,0.04,0.05,0.06,0.07,0.08,0.09,0.10]
# F_list_noise = 
#  [0.97062911 0.97062911 0.96996106 0.9679874  0.96472136 0.96018506 0.95440917 0.94743255 0.93930171 0.93007021 0.919798   0.90855067]

#  time_cost is:  182258.63508796692


#drift
# noise_am_list = [-0.020,-0.018,-0.016,-0.014,-0.012,-0.010,-0.008,-0.006,-0.004,-0.002,0.000,0.002,0.004,0.006,0.008,0.010,0.012,0.014,0.016,0.018,0.020]

F_list_noise = 
 [0.97062911 0.96608198 0.96696784 0.9677596  0.96845644 0.96905766 0.96956262 0.96997081 0.97028182 0.97049533 0.97061113 0.97062911 0.97054928 0.97037175 0.97009673 0.96972452 0.96925555 0.96869033 0.96802949 0.96727375 0.96642391 0.96548089]

 time_cost is:  189040.30624771118

